from .vit import VisionTransformer as vvit
from .single_bb import (
    vmoe,
    vease,
    vbase_single,
    vbase_multi,
)
from .multi_bb import (
    vmulti,
    vmulti_wbb,
    vmulti_sn,
    vmulti_dsn,
    vmulti_unipt,
)


__all__ = [
    "vvit",
    "vmoe",
    "vease",
    "vbase_single",
    "vbase_multi",
    "vmulti",
    "vmulti_wbb",
    "vmulti_sn",
    "vmulti_dsn",
    "vmulti_unipt",
]

import timm
from utils.toolkit import NamespaceDict


def get_tuning_config(args, extra_params=None):
    # FIXME: too many config mixed in args
    config = {
        "embed_dim": 768,
        "peft_name": args.get("peft_name", None),
        "ffn_option": args.get("ffn_option", "none"),
        "ffn_rank": args.get("ffn_rank", None),
        "ffn_adapter_scalar": args.get("ffn_adapter_scalar", 1.0),
        "ffn_adapter_layernorm_option": args.get(
            "ffn_adapter_layernorm_option", "none"
        ),
        "mnc_norm": args.get("mnc_norm", False),
        "moe_topk": args.get("moe_topk", None),
        "cam_visual": args["cam_visual"],
        "_device": args["device"][0],
    }
    config.update({"sidenet_cfg": args.get("sidenet_cfg", None)})
    if extra_params:
        config.update(extra_params)
    return NamespaceDict(config)


def get_backbone(
    args,
):
    name = args["backbone_type"]

    config = get_tuning_config(args)
    model = timm.create_model(
        name,
        num_classes=0,
        global_pool=False,
        drop_path_rate=0.0,
        config=config,
    )
    model.out_dim = 768
    return model.eval()
